import rclpy
from rclpy.node import Node
import numpy as np
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped, TransformStamped
from tf2_ros import TransformListener
from tf2_ros.buffer import Buffer
import tf2_ros
from rclpy.time import Time
import cv2
from nav_msgs.msg import OccupancyGrid,Odometry
import heapq
import math

class Map2Map(Node):

    def __init__(self):
        super().__init__('map_2_map_node')

        # xy搜索范围（单位：网格数）
        self.declare_parameter('xy_search_range', 20)
        # 角度搜索范围（单位：度）
        self.declare_parameter('theta_search_range', 90)
        # xy搜索步长（单位：网格数）
        self.declare_parameter('xy_search_dt', 1)
        # 角度搜索步长（单位：度）
        self.declare_parameter('theta_search_dt', 90)
        
        # 全局地图
        self.global_cloud_sub = self.create_subscription(
            OccupancyGrid, 
            '/inflated_map',
            self.global_cloud_callback ,
            10)
        self.global_cloud_sub
        
        # 点云地图
        self.local_cloud_sub = self.create_subscription(
            OccupancyGrid, 
            '/raw_costmap',
            self.local_cloud_callback ,
            10)
        self.local_cloud_sub

        # 里程计
        self.odom_sub = self.create_subscription(
            Odometry,
            '/Odometry',
            self.odom_callback,
            10)
        self.odom_sub
        
        # 定时器
        timer_period = 0.01  # 20ms = 0.02s
        self.timer = self.create_timer(timer_period, self.timer_callback)
        
        self.tf_broadcaster = tf2_ros.TransformBroadcaster(self)
        
        self.gloabal_map_msg=None
        self.base_link_pose=None
        self.odom_on_map_x_best=0.0
        self.odom_on_map_y_best=0.0
        self.best_theta=0.0
        self.call_in_tf=0
        self.max_matching_value=0
        
    def global_cloud_callback(self,msg):
        if not self.gloabal_map_msg:
            self.gloabal_map_msg=msg

    def odom_callback(self,msg):
        self.base_link_pose=msg.pose
        
    def local_cloud_callback(self,msg):
        if not self.gloabal_map_msg or not self.base_link_pose:
            return
        # 裁剪地图
        if self.call_in_tf>50:
            return

        local_width=msg.info.width
        local_height=msg.info.height
        local_origin_x=msg.info.origin.position.x
        local_origin_y=msg.info.origin.position.y
        local_map_array=np.array(msg.data).reshape((local_height, local_width))
        
        global_width=self.gloabal_map_msg.info.width
        global_height=self.gloabal_map_msg.info.height
        global_origin_x=self.gloabal_map_msg.info.origin.position.x
        global_origin_y=self.gloabal_map_msg.info.origin.position.y
        gloabal_map_array=np.array(self.gloabal_map_msg.data).reshape((global_height, global_width))

        xy_search_dt_value=self.get_parameter('xy_search_dt').get_parameter_value().integer_value
        xy_search_range_value=self.get_parameter('xy_search_range').get_parameter_value().integer_value

        for i in range(-xy_search_range_value,xy_search_range_value,xy_search_dt_value):
            for j in range(-xy_search_range_value,xy_search_range_value,xy_search_dt_value):
                base_on_map_x=local_origin_x+i*self.gloabal_map_msg.info.resolution
                base_on_map_y=local_origin_y+j*self.gloabal_map_msg.info.resolution
                local_in_global_x=base_on_map_x-global_origin_x
                local_in_global_y=base_on_map_y-global_origin_y

                if local_in_global_x<0 or local_in_global_y<0:
                    self.get_logger().info("local_costmap out of global_costmap")
                    continue

                local_in_globalmap_x,local_in_globalmap_y=self.world_to_map(local_in_global_x,local_in_global_y)
                if local_in_globalmap_x+local_width>global_width or local_in_globalmap_y+local_height>global_height:
                    # self.get_logger().info("local_costmap out of global_costmap")
                    continue
                
                
                current_value=0
                for x in range(0,local_width-1):
                    for y in range(0,local_height-1):
                        if local_map_array[y][x]==gloabal_map_array[y+local_in_globalmap_y][x+local_in_globalmap_x]:
                            current_value+=1
                
                if current_value>self.max_matching_value:
                    self.max_matching_value=current_value
                    self.odom_on_map_y_best=local_height/2*self.gloabal_map_msg.info.resolution+base_on_map_y-self.base_link_pose.pose.position.y
                    self.odom_on_map_x_best=local_height/2*self.gloabal_map_msg.info.resolution+base_on_map_x-self.base_link_pose.pose.position.x
            self.call_in_tf+=1
        

    # 坐标发布

    def timer_callback(self):
        transform_stamped = TransformStamped()
        transform_stamped.header.stamp = self.get_clock().now().to_msg()
        transform_stamped.header.frame_id = 'map'
        transform_stamped.child_frame_id = 'odom'

        # 设置 map 到 odom 的变换关系，假设没有变换
        transform_stamped.transform.translation.x = self.odom_on_map_x_best
        transform_stamped.transform.translation.y = self.odom_on_map_y_best
        # self.get_logger().info(f'gx: {self.odom_on_map_x_best}')
        # self.get_logger().info(f'gy: {self.odom_on_map_y_best}')
        transform_stamped.transform.translation.z = 0.0

        quaternion = self.euler_to_quaternion(0,0,self.best_theta)
        transform_stamped.transform.rotation.x = quaternion[0]
        transform_stamped.transform.rotation.y = quaternion[1]
        transform_stamped.transform.rotation.z = quaternion[2]
        transform_stamped.transform.rotation.w = quaternion[3]
        self.tf_broadcaster.sendTransform(transform_stamped)

    # 世界坐标系转换到地图坐标系
    def world_to_map(self, x, y):
        mx = int(x / self.gloabal_map_msg.info.resolution)
        my = int(y / self.gloabal_map_msg.info.resolution)
        return mx, my
    
    # 欧拉角转四元数
    def euler_to_quaternion(self,roll, pitch, yaw):
        cy = np.cos(yaw * 0.5)
        sy = np.sin(yaw * 0.5)
        cp = np.cos(pitch * 0.5)
        sp = np.sin(pitch * 0.5)
        cr = np.cos(roll * 0.5)
        sr = np.sin(roll * 0.5)

        qw = cy * cp * cr + sy * sp * sr
        qx = cy * cp * sr - sy * sp * cr
        qy = sy * cp * sr + cy * sp * cr
        qz = sy * cp * cr - cy * sp * sr

        return np.array([qx, qy, qz, qw])


def main(args=None):
    rclpy.init(args=args)

    map_2_map_node = Map2Map()

    rclpy.spin(map_2_map_node)

    map_2_map_node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
